package com.peng.leetcode.unionFind;

import java.util.*;

/**
 * AccountMerge
 *
 * @author: lupeng6
 * @create: 2021/1/18 11:18
 */
public class AccountMerge {

    public static void main(String[] args) {
        String[][] accounts = {{"John", "johnsmith@mail.com", "john00@mail.com"}, {"John", "johnnybravo@mail.com"}, {"John", "johnsmith@mail.com", "john_newyork@mail.com"}, {"Mary", "mary@mail.com"}};
        List<List<String>> list = new ArrayList<>();
        for (String[] account : accounts) {
            list.add(new ArrayList(Arrays.asList(account)));
        }
        List<List<String>> lists = new AccountMerge().accountsMerge(list);
    }

    public List<List<String>> accountsMerge(List<List<String>> accounts) {
        Map<String, Integer> mailToIndex = new HashMap<>();
        Map<String, String> mailToName = new HashMap<>();

        for (int i = 0; i < accounts.size(); i++) {
            List<String> account =  accounts.get(i);
            String name = account.get(0);

            for (int j = 1; j < account.size(); j++) {
                String mail = account.get(j);
                mailToName.put(name, mail);
                mailToIndex.put(mail, i);
            }
        }

        UnionFind unionFind = new UnionFind(accounts.size());

        // 将邮箱索引合并
        for (int i = 0; i < accounts.size(); i++) {
            List<String> account = accounts.get(i);
            String firstMail = account.get(1);
            Integer firstMailIndex = mailToIndex.get(firstMail);
            for (int j = 2; j < account.size(); j++) {
                String nextMail = account.get(j);
                Integer nextMailIndex = mailToIndex.get(nextMail);
                unionFind.union(firstMailIndex, nextMailIndex);
            }
        }

        Map<Integer, List<String>> indexToMails = new HashMap<>();

        for (String mail : mailToIndex.keySet()) {
            Integer index = mailToIndex.get(mail);
            int rootIndex = unionFind.find(index);
            indexToMails.putIfAbsent(rootIndex, new ArrayList<>());
            indexToMails.get(rootIndex).add(mail);
        }

        List<List<String>> result = new ArrayList<>();

        for (Integer index : indexToMails.keySet()) {
            List<String> item = new ArrayList<>();

            String name = accounts.get(index).get(0);
            List<String> mails = indexToMails.get(index);
            Collections.sort(mails);

            item.add(name);
            item.addAll(mails);

            result.add(item);
        }

        return result;
    }

    static class UnionFind {
        List<Integer> parent = null;

        public UnionFind(int size) {
            parent = new ArrayList<>(size);
            for (int i = 0; i < size; i++) {
                parent.add(i);
            }
        }

        public int find(int x) {
            Integer xRoot = parent.get(x);

            if (xRoot != x) {
                parent.set(x, find(xRoot));
                x = parent.get(x);
            }
            return x;
        }

        public void union(int x, int y) {
            int xRoot = find(x);
            int yRoot = find(y);

            if (xRoot != yRoot) {
                parent.set(xRoot, yRoot);
            }
        }
    }
}
